
#include <stdio.h>
#include <errno.h>
#include <cstring>
#include <string>
extern "C"
{
#include <libavutil/hwcontext.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
}

static AVFormatContext *ifmt_ctx = NULL, *ofmt_ctx = NULL;
static AVBufferRef *hw_device_ctx = NULL;
static AVCodecContext *decoder_ctx = NULL, *encoder_ctx = NULL;
static int video_stream = -1;
static AVStream *ost;
static int initialized = 0;

static enum AVPixelFormat GetQsvFormat(AVCodecContext * codec_ctx, const enum AVPixelFormat * pix_fmts)
{
    const enum AVPixelFormat * p;
    for (p = pix_fmts; *p != AV_PIX_FMT_NONE; p++)
    {
        if (*p == AV_PIX_FMT_QSV)
        {
            return AV_PIX_FMT_QSV;
        }
    }
    return AV_PIX_FMT_NONE;
}

static enum AVPixelFormat GetVaApiFormat(AVCodecContext * codec_ctx, const enum AVPixelFormat * pix_fmts)
{
    const enum AVPixelFormat * p;
    for (p = pix_fmts; *p != AV_PIX_FMT_NONE; p++)
    {
        if (*p == AV_PIX_FMT_VAAPI)
        {
            return AV_PIX_FMT_VAAPI;
        }
    }
    return AV_PIX_FMT_NONE;
}

static int OpenInputFile(const std::string & filename)
{
    int result = 0;
    result = avformat_open_input(&ifmt_ctx, filename.c_str(), NULL, NULL);
    if (result < 0)
    {
        printf("avformat_open_input.\n");
        return -1;
    }
    avformat_find_stream_info(ifmt_ctx, nullptr);

    const AVCodec * decoder = nullptr;
    result = av_find_best_stream(ifmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
    if (result < 0)
    {
        printf("av_find_best_stream.\n");
        return -1;
    }
    decoder = avcodec_find_decoder_by_name("h264_qsv");
    // video stream index
    video_stream = result;

    decoder_ctx = avcodec_alloc_context3(decoder);
    if (!decoder_ctx)
    {
        printf("avcodec_alloc_context3.\n");
        return -1;
    }
    AVStream * video = ifmt_ctx->streams[video_stream];
    avcodec_parameters_to_context(decoder_ctx, video->codecpar);
    decoder_ctx->pkt_timebase = video->time_base;
    decoder_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
    decoder_ctx->get_format = GetQsvFormat;
    decoder_ctx->codec_id   = AV_CODEC_ID_H264;

    result = avcodec_open2(decoder_ctx, decoder, nullptr);
    if (result < 0)
    {
        printf("avcodec_open2.\n");
        return -1;
    }
    av_dump_format(ifmt_ctx, 0, filename.c_str(), 0);
    return 0;
}

static int encode_write(AVPacket * packet, AVFrame * frame)
{
    int result = 0;
    result = avcodec_send_frame(encoder_ctx, frame);
    if (result < 0)
    {
        goto fail;
    }
    while (1)
    {
        result = avcodec_receive_packet(encoder_ctx, packet);
        if (result)
        {
            break;
        }
        packet->stream_index = 0;
        av_packet_rescale_ts(packet, ifmt_ctx->streams[video_stream]->time_base, ofmt_ctx->streams[0]->time_base);
        result = av_interleaved_write_frame(ofmt_ctx, packet);
        if (result < 0)
        {
            return -1;
        }
    }
fail:
    if (result == AVERROR_EOF)
    {
        return 0;
    }
    result = result == AVERROR(EAGAIN) ? 0 : -1;
    return result;
}

static int dec_enc(AVPacket * packet, const AVCodec * enc_dec)
{
    AVFrame * frame;

    int result = avcodec_send_packet(decoder_ctx, packet);
    if (result < 0)
    {
        char buffer[1024] = {0};
        av_make_error_string(buffer, 1000, result);
        printf("avcodec_send_packet %s.\n", buffer);
        return -1;
    }
    while (1)
    {
        if (!(frame = av_frame_alloc()))
        {
            return -1;
        }
        result = avcodec_receive_frame(decoder_ctx, frame);
        if (result == AVERROR(EAGAIN) || result == AVERROR_EOF)
        {
            av_frame_free(&frame);
            return 0;
        }
        else if (result < 0)
        {
            av_frame_free(&frame);
            return -1;
        }
        if (!initialized)
        {
            encoder_ctx->hw_frames_ctx = av_buffer_ref(decoder_ctx->hw_frames_ctx);
            if (!encoder_ctx->hw_frames_ctx)
            {
                result = AVERROR(ENOMEM);
                goto fail;
            }

            avcodec_parameters_to_context(encoder_ctx, ifmt_ctx->streams[video_stream]->codecpar);

            encoder_ctx->framerate = ifmt_ctx->streams[video_stream]->avg_frame_rate;
            encoder_ctx->time_base = ifmt_ctx->streams[video_stream]->time_base;
            encoder_ctx->pix_fmt   = AV_PIX_FMT_QSV;

            if ((result = avcodec_open2(encoder_ctx, enc_dec, NULL)) < 0)
            {
                goto fail;
            }

            if (!(ost = avformat_new_stream(ofmt_ctx, NULL)))
            {
                result = AVERROR(ENOMEM);
                goto fail;
            }

            ost->time_base = encoder_ctx->time_base;
            result = avcodec_parameters_from_context(ost->codecpar, encoder_ctx);
            if (result < 0)
            {
                goto fail;
            }

            /* write the stream header */
            if ((result = avformat_write_header(ofmt_ctx, NULL)) < 0)
            {
                goto fail;
            }
            initialized = 1;
        }
        if ((result = encode_write(packet, frame)) < 0)
        {
        }

        fail:
        av_frame_free(&frame);
        if (result < 0)
            return result;
    }
    return 0;
}

int main(int argc, char **argv)
{
    const AVCodec *enc_codec;
    int ret = 0;
    AVPacket *dec_pkt;

    if (argc != 3)
    {
        fprintf(stderr, "Usage: %s <input file> <encode codec> <output file>\n"
                        "The output format is guessed according to the file extension.\n"
                        "\n", argv[0]);
        return -1;
    }

    ret = av_hwdevice_ctx_create(&hw_device_ctx, AV_HWDEVICE_TYPE_QSV, "auto", NULL, 0);
    if (ret < 0)
    {
        printf("av_hwdevice_ctx_create.\n");
        return -1;
    }

    dec_pkt = av_packet_alloc();
    if (!dec_pkt)
    {
        printf("av_packet_alloc.\n");
        goto end;
    }

    if ((ret = OpenInputFile(argv[1])) < 0)
        goto end;

    if (!(enc_codec = avcodec_find_encoder_by_name("h264_qsv")))
    {
        printf("avcodec_find_encoder.\n");
        ret = -1;
        goto end;
    }

    if ((ret = (avformat_alloc_output_context2(&ofmt_ctx, NULL, NULL, argv[2]))) < 0)
    {
        printf("avformat_alloc_output_context2.\n");
        goto end;
    }

    if (!(encoder_ctx = avcodec_alloc_context3(enc_codec)))
    {
        printf("avcodec_alloc_context3.\n");
        ret = AVERROR(ENOMEM);
        goto end;
    }

    ret = avio_open(&ofmt_ctx->pb, argv[2], AVIO_FLAG_WRITE);
    if (ret < 0)
    {
        printf("avio_open.\n");
        goto end;
    }

    /* read all packets and only transcoding video */
    while (ret >= 0)
    {
        if ((ret = av_read_frame(ifmt_ctx, dec_pkt)) < 0)
            break;

        if (video_stream == dec_pkt->stream_index)
            ret = dec_enc(dec_pkt, enc_codec);

        av_packet_unref(dec_pkt);
    }

    /* flush decoder */
    av_packet_unref(dec_pkt);
    ret = dec_enc(dec_pkt, enc_codec);

    /* flush encoder */
    ret = encode_write(dec_pkt, NULL);

    /* write the trailer for output stream */
    av_write_trailer(ofmt_ctx);

    end:
    avformat_close_input(&ifmt_ctx);
    avformat_close_input(&ofmt_ctx);
    avcodec_free_context(&decoder_ctx);
    avcodec_free_context(&encoder_ctx);
    av_buffer_unref(&hw_device_ctx);
    av_packet_free(&dec_pkt);
    return ret;
}



